import os
import subprocess
import shutil
import time

def show_welcome_message():
    """显示欢迎信息"""
    banner = """
    *********************************************************
    *                                                       *
    *        欢迎使用镜像站选择工具                           *
    *                                                       *
    * 本工具旨在将：                                         *
    * /etc/ros/rosdep/sources.list.d/20-default.list 中的地址*
    * 更换为国内镜像站的地址                                  *
    *********************************************************    

    """
    print(banner)
    print("请从下方选择要使用的镜像站：\n")

def show_mirror_options():
    """显示镜像站选项"""
    options = {
        "1": "清华大学开源镜像站",
        "2": "北京外国语大学镜像站",
        "3": "中国科学技术大学开源镜像站",
        "4": "南京大学开源镜像站",
        "5": "退出程序"
    }
    
    for key, value in options.items():
        print(f"[{key}] {value}")
    
    return options

def check_curl_installed():
    """检查系统中是否安装了curl"""
    return shutil.which("curl") is not None

def install_curl():
    """尝试安装curl"""
    print("检测到系统中未安装curl，尝试安装...")
    try:
        subprocess.run(["sudo", "apt-get", "update"], check=True)
        subprocess.run(["sudo", "apt-get", "install", "-y", "curl"], check=True)
        print("curl安装成功！")
        return True
    except subprocess.CalledProcessError as e:
        print(f"安装curl失败: {e}")
        return False

def execute_tsinghua_commands():
    """执行清华大学镜像站的相关命令"""
    if not check_curl_installed():
        if not install_curl():
            print("无法继续执行，因为curl安装失败")
            return False
    
    commands = [
        "sudo mkdir -p /etc/ros/rosdep/sources.list.d/",
        "sudo curl -o /etc/ros/rosdep/sources.list.d/20-default.list -L https://mirrors.tuna.tsinghua.edu.cn/github-raw/ros/rosdistro/master/rosdep/sources.list.d/20-default.list",
        "export ROSDISTRO_INDEX_URL=https://mirrors.tuna.tsinghua.edu.cn/rosdistro/index-v4.yaml",
        "echo 'export ROSDISTRO_INDEX_URL=https://mirrors.tuna.tsinghua.edu.cn/rosdistro/index-v4.yaml' >> ~/.bashrc",
        "echo '环境变量已添加到 ~/.bashrc，请打开新终端执行 source ~/.bashrc使设置生效'"
    ]
    
    print("\n正在执行清华大学开源镜像站设置命令...")
    return execute_commands(commands, "清华大学")

def execute_bfsu_commands():
    """执行北京外国语大学镜像站的相关命令"""
    if not check_curl_installed():
        if not install_curl():
            print("无法继续执行，因为curl安装失败")
            return False
    
    commands = [
        "sudo mkdir -p /etc/ros/rosdep/sources.list.d/",
        "sudo curl -o /etc/ros/rosdep/sources.list.d/20-default.list -L https://mirrors.bfsu.edu.cn/github-raw/ros/rosdistro/master/rosdep/sources.list.d/20-default.list",
        "export ROSDISTRO_INDEX_URL=https://mirrors.bfsu.edu.cn/rosdistro/index-v4.yaml",
        "echo 'export ROSDISTRO_INDEX_URL=https://mirrors.bfsu.edu.cn/rosdistro/index-v4.yaml' >> ~/.bashrc",
        "echo '环境变量已添加到 ~/.bashrc，请打开新终端执行 source ~/.bashrc使设置生效'"
    ]
    
    print("\n正在执行北京外国语大学镜像站设置命令...")
    return execute_commands(commands, "北京外国语大学")

def execute_ustc_commands():
    """执行中国科学技术大学镜像站的相关命令"""
    if not check_curl_installed():
        if not install_curl():
            print("无法继续执行，因为curl安装失败")
            return False
    
    commands = [
        "sudo mkdir -p /etc/ros/rosdep/sources.list.d/",
        "sudo curl -o /etc/ros/rosdep/sources.list.d/20-default.list https://mirrors.ustc.edu.cn/rosdistro/rosdep/sources.list.d/20-default.list",
        "sudo sed -i 's#raw.githubusercontent.com/ros/rosdistro/master#mirrors.ustc.edu.cn/rosdistro#g' /etc/ros/rosdep/sources.list.d/20-default.list",
        "export ROSDISTRO_INDEX_URL=https://mirrors.ustc.edu.cn/rosdistro/index-v4.yaml",
        "echo 'export ROSDISTRO_INDEX_URL=https://mirrors.ustc.edu.cn/rosdistro/index-v4.yaml' >> ~/.bashrc",
        "echo '环境变量已添加到 ~/.bashrc，请打开新终端执行 source ~/.bashrc使设置生效'"
    ]
    
    print("\n正在执行中国科学技术大学镜像站设置命令...")
    return execute_commands(commands, "中国科学技术大学")

def execute_nju_commands():
    """执行南京大学镜像站的相关命令"""
    if not check_curl_installed():
        if not install_curl():
            print("无法继续执行，因为curl安装失败")
            return False
    
    commands = [
        "sudo mkdir -p /etc/ros/rosdep/sources.list.d/",
        "sudo curl -o /etc/ros/rosdep/sources.list.d/20-default.list -L https://mirrors.cernet.edu.cn/github-raw/ros/rosdistro/master/rosdep/sources.list.d/20-default.list",
        "export ROSDISTRO_INDEX_URL=https://mirror.nju.edu.cn/rosdistro/index-v4.yaml",
        "echo 'export ROSDISTRO_INDEX_URL=https://mirror.nju.edu.cn/rosdistro/index-v4.yaml' >> ~/.bashrc",
        "echo '环境变量已添加到 ~/.bashrc，请打开新终端执行 source ~/.bashrc使设置生效'"
    ]
    
    print("\n正在执行南京大学开源镜像站设置命令...")
    return execute_commands(commands, "南京大学")

def execute_commands(commands, mirror_name):
    """通用命令执行函数"""
    success = True
    for cmd in commands:
        print(f"\n执行命令: {cmd}")
        try:
            result = subprocess.run(cmd, shell=True, check=True, 
                                  stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                  text=True)
            print(result.stdout)
        except subprocess.CalledProcessError as e:
            print(f"命令执行失败: {e}")
            print(f"错误输出: {e.stderr}")
            success = False
        except Exception as e:
            print(f"执行命令时发生意外错误: {e}")
            success = False
    
    if success:
        print(f"\n{mirror_name}镜像站配置文件设置完成！")
    else:
        print(f"\n{mirror_name}镜像站配置文件设置过程中出现错误！")
    
    return success

def run_rosdep_update(mirror_name):
    """执行rosdep update命令"""
    print("\n" + "="*50)
    print("正在执行 rosdep update...")
    print("="*50)
    
    try:
        result = subprocess.run("rosdep update", shell=True, check=True,
                              stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                              text=True)
        print(result.stdout)
        print(f"\n✓ 已经使用{mirror_name}更新完成！")
        print(f"程序将在5秒后自动退出...")
        
        # 5秒倒计时
        for i in range(5, 0, -1):
            print(f"{i}秒后退出...", end='\r')
            time.sleep(1)
        
        print("\n感谢使用！")
        return True
    except subprocess.CalledProcessError as e:
        print(f"rosdep update 执行失败: {e}")
        print(f"错误输出: {e.stderr}")
        print("\n提示: rosdep命令执行失败，请确保已安装rosdep")
        print("可以尝试运行: sudo apt-get install python3-rosdep")
        return False
    except Exception as e:
        print(f"执行rosdep update时发生意外错误: {e}")
        return False

def main():
    """主程序入口"""
    show_welcome_message()
    
    while True:
        options = show_mirror_options()
        choice = input("\n请输入您的选择：")
        
        if choice == "5":
            print("\n感谢使用，再见！")
            break
        elif choice == "1":
            print(f"\n您选择了: {options[choice]}")
            if execute_tsinghua_commands():
                if run_rosdep_update("清华大学开源镜像站"):
                    break  # 成功后退出程序
        elif choice == "2":
            print(f"\n您选择了: {options[choice]}")
            if execute_bfsu_commands():
                if run_rosdep_update("北京外国语大学镜像站"):
                    break  # 成功后退出程序
        elif choice == "3":
            print(f"\n您选择了: {options[choice]}")
            if execute_ustc_commands():
                if run_rosdep_update("中国科学技术大学开源镜像站"):
                    break  # 成功后退出程序
        elif choice == "4":
            print(f"\n您选择了: {options[choice]}")
            if execute_nju_commands():
                if run_rosdep_update("南京大学开源镜像站"):
                    break  # 成功后退出程序
        else:
            print("\n无效输入，请重新选择！")

if __name__ == '__main__':
    main()